#include "xpelabmd5sum.h"
#include "sm4encryptmgr.h"


#define MD5_SIZE 16
#define CHK_HEAD_LGH 8
static const char CHK_FILE_HEAD[CHK_HEAD_LGH] = {0x3c, 0x53, 0x63, 0x72,0x69, 0x70, 0x74, 0x20};   //"<Script " by UTF-8

namespace A {
namespace B {


    unsigned char SM4EncryptMgr::mGenerateKey[SM4_BLOCK_SIZE] = {
        0x01, 0xd3, 0x45, 0x67, 0xaa, 0xab, 0xcd, 0xef,
        0xf0, 0xcc, 0xba, 0x98, 0x76, 0x51, 0xdf, 0x1d
        };

    long SM4EncryptMgr::resizeEncryptLength(long length)
    {
        if(length % SM4_BLOCK_SIZE)
        {
            return (length/SM4_BLOCK_SIZE + 1)*SM4_BLOCK_SIZE;
        }

        return length;
    }

    bool SM4EncryptMgr::setEncryptFileBuf(unsigned char* buf, FILE* file, long filelength, long buflength)
    {
        memset(buf, 0, filelength);

        int a = fread(buf, filelength, 1, file);
        if(a != 1)
        {
            return false;
        }

        //tail 16 bit align by 0xff
        if(buflength > filelength)
        {
           memset(buf+filelength, 0xFF, buflength-filelength);
        }

        return true;
    }

    bool SM4EncryptMgr::encryptStr(const string& in, string& res)
    {
        res.clear();
        string orgStr(in);
        if(in.empty())
        {
            return false;
        }

        long length = in.length();
        long reslen = resizeEncryptLength(length);
        unsigned char* resBuf = new unsigned char[reslen];

        orgStr.append(reslen-length, 0xFF);
        
        if(!encryptData((unsigned char*)&orgStr[0], resBuf, reslen))
        {
            delete [] resBuf;
            return false;
        }

        res.assign((char*)resBuf, reslen);

        delete [] resBuf;
        return true;

    }

    bool SM4EncryptMgr::encryptFile(const string& infilename, const string& outfilename)
    {

        if (!xPeLabFiles::exist(infilename))
        {
            return false;
        }

        FILE *file = fopen(infilename.c_str(), "rb");
        if(!file)
        {
            return false;
        }

        fseek(file, 0L, SEEK_END);
        long length= ftell(file);
        fseek(file, 0L, SEEK_SET);

        if (length < 0)
        {
            return false;
        }

        int buflength = resizeEncryptLength(length);
        unsigned char* buf = new unsigned char[buflength];

        if(!setEncryptFileBuf(buf, file, length, buflength))
        {
            delete [] buf;
            fclose(file);
            return false;
        }
        fclose(file);


        if(!encryptData(buf, buf, buflength))
        {
            delete [] buf;
            return false;
        }

        FILE *outfile = fopen(outfilename.c_str(), "wb");
        if(!outfile)
        {
            delete [] buf;
            return false;
        }

        //add chk sequnce head
        if(!addMD5SumHead(outfile, buf, buflength))
        {
            fclose(outfile);
            delete [] buf;
            return false;
        }

        //write context
        if(fwrite(buf, buflength, 1, outfile) != 1)
        {
            fclose(outfile);
            delete [] buf;
            return false;
        }

        fclose(outfile);
        delete [] buf;

        return true;

    }

    bool SM4EncryptMgr::encryptData(const unsigned char* in, unsigned char* out, long buflength)
    {
        if(buflength % SM4_BLOCK_SIZE)
        {
            return false;
        }

        int cycle = buflength / SM4_BLOCK_SIZE;
        SM4_KEY key;
        m_encryptProto.sm4_set_key(mGenerateKey, &key);

        unsigned char block[SM4_BLOCK_SIZE] = {0};
        for(int i = 0; i < cycle; i ++)
        {
            memcpy(block, (in + i*SM4_BLOCK_SIZE), SM4_BLOCK_SIZE);

            m_encryptProto.sm4_encrypt(block, block, &key);

            memcpy((out + i*SM4_BLOCK_SIZE), block, SM4_BLOCK_SIZE);
        }

        return true;

    }

    bool SM4EncryptMgr::dencryptStr(const string& in, string& res)
    {
        res.clear();
        if(in.empty())
        {
            return false;
        }

        long orgLen = in.length();
        long resLen = 0;
        unsigned char* resBuf = new unsigned char[orgLen];

        if(!decryptData((unsigned char*)&in[0], resBuf, orgLen, resLen))
        {
            delete [] resBuf;
            return false;
        }

        res.assign((char*)resBuf, resLen);

        delete [] resBuf;
        return true;
        
    }

    bool SM4EncryptMgr::decryptFile(const string& infilename, const string& outfilename)
    {
        if (!xPeLabFiles::exist(infilename))
        {
            return false;
        }

        FILE *file = fopen(infilename.c_str(), "rb");
        if(!file)
        {
            return false;
        }

        fseek(file, 0L, SEEK_END);
        long length= ftell(file);
        fseek(file, 0L, SEEK_SET);

        if (length < 0 )
        {
            return false;
        }

        unsigned char* buf = new unsigned char[length];
        memset(buf, 0, length);
        if(fread(buf, length, 1, file) != 1)
        {
            delete[] buf;
            fclose(file);
            return false;
        }
        fclose(file);

        if(!chkMD5(buf, length))
        {
            delete[] buf;
            return false;
        }

        unsigned char *context = buf+MD5_SIZE;
        long filelength = -1;
        if(!decryptData(context, context, length-MD5_SIZE,filelength))
        {
            delete [] buf;
            return false;
        }


        FILE *outfile = fopen(outfilename.c_str(), "wb");
        if(!outfile)
        {
            delete [] buf;
            return false;
        }

        if(fwrite(context, filelength, 1, outfile) != 1)
        {
            fclose(outfile);
            delete [] buf;
            return false;
        }

        fclose(outfile);
        delete [] buf;

        return true;

    }

    char* SM4EncryptMgr::getDecryptDataFromFile(const string& infilename, long& filelength)
    {
        if (!xPeLabFiles::exist(infilename))
        {
            return nullptr;
        }

        FILE *file = fopen(infilename.c_str(), "rb");
        if(!file)
        {
            return nullptr;
        }

        fseek(file, 0L, SEEK_END);
        long length= ftell(file);
        fseek(file, 0L, SEEK_SET);

        if (length < 0 )
        {
            return nullptr;
        }

        unsigned char* buf = new unsigned char[length];
        memset(buf, 0, length);
        if(fread(buf, length, 1, file) != 1)
        {
            delete[] buf;
            fclose(file);
            return nullptr;
        }
        fclose(file);

        if(!chkMD5(buf, length))
        {
            delete[] buf;
            return nullptr;
        }

        unsigned char *context = buf+MD5_SIZE;
        if(!decryptData(context, context, length-MD5_SIZE,filelength))
        {
            delete [] buf;
            return nullptr;
        }

        char *resData = new char[filelength];
        memcpy(resData, context, filelength);

        delete [] buf;
        return resData;

    }

    bool SM4EncryptMgr::decryptData(const unsigned char* in, unsigned char* out, long buflength, long& filelength)
    {
        if(buflength % SM4_BLOCK_SIZE)
        {
            return false;
        }

        SM4_KEY key;
        m_encryptProto.sm4_set_key(mGenerateKey, &key);

        int cycle = buflength / SM4_BLOCK_SIZE;

        unsigned char block[SM4_BLOCK_SIZE] = {0};
        for(int i = 0; i < cycle; i ++)
        {
            memcpy(block, (in + i*SM4_BLOCK_SIZE), SM4_BLOCK_SIZE);

            m_encryptProto.sm4_decrypt(block, block, &key);

            if(i == (cycle-1))
            {
                int j = 0;
                while( (j < SM4_BLOCK_SIZE) && block[j] != 0xff )
                {
                    j++;
                }
                memcpy((out + i*SM4_BLOCK_SIZE), block, j);
                filelength = i*SM4_BLOCK_SIZE + j;
            }
            else
                memcpy((out + i*SM4_BLOCK_SIZE), block, SM4_BLOCK_SIZE);
        }


        return true;
    }

    bool SM4EncryptMgr::isEncryptFile(const std::string& filename)
    {
        FILE *file = fopen(filename.c_str(), "rb");
        if(!file)
        {
            return false;
        }

        char filehead[CHK_HEAD_LGH] = {0};
        if(fread(filehead, CHK_HEAD_LGH, 1, file) != 1)
        {
            fclose(file);
            return false;
        }
        fclose(file);

        return (strncmp(filehead, CHK_FILE_HEAD, CHK_HEAD_LGH) != 0);

    }

    bool SM4EncryptMgr::addMD5SumHead(FILE *file, unsigned char *buf, long buflength)
    {
        /*
         * MD5 16 bits sequence
         * --------------
         * encrypt 16 bits group context
         * --------------
         * tail 16 bits align 0xff
         *---------------
         */
        unsigned char md5value[MD5_SIZE];
        xPeLabMD5Sum::md5sumvalue(buf, buflength, md5value);

        if(fwrite(md5value, MD5_SIZE, 1, file) != 1)
        {
            return false;
        }
        return true;
    }

    bool SM4EncryptMgr::chkMD5(const unsigned char *filebuf, long filelength)
    {
        if(!filebuf && (filelength < MD5_SIZE))
            return false;

        unsigned char md5value[MD5_SIZE];
        xPeLabMD5Sum::md5sumvalue(filebuf + MD5_SIZE, filelength - MD5_SIZE, md5value);

        for (int i = 0; i < MD5_SIZE; i++)
        {
            if(filebuf[i] != md5value[i])
                return false;
        }

        return true;
    }
}
}
